#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
重置 LiteLLM 配置
"""

import os
import sys
import json

# 添加项目根目录到路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from app.core.config import settings
import litellm


def reset_litellm():
    """重置 LiteLLM 配置"""
    print("重置 LiteLLM 配置...")
    
    # 打印当前配置
    print(f"当前 LiteLLM 配置:")
    print(f"- API 密钥: {litellm.api_key}")
    print(f"- Ollama API 基础 URL: {litellm.ollama_api_base}")
    print(f"- 模型列表: {json.dumps(litellm.model_list, indent=2) if hasattr(litellm, 'model_list') else '无'}")
    
    # 重置配置
    litellm.api_key = settings.LITELLM_API_KEY
    litellm.ollama_api_base = settings.OLLAMA_BASE_URL
    
    # 设置模型配置
    if settings.LITELLM_MODEL.startswith("ollama/"):
        model_name = settings.LITELLM_MODEL.split("/")[1]  # 获取模型名称，例如 "qwen2.5:latest"
        litellm.model_list = [
            {
                "model_name": f"ollama/{model_name}",
                "litellm_params": {
                    "api_base": settings.OLLAMA_BASE_URL
                }
            }
        ]
    
    # 设置环境变量
    os.environ["OLLAMA_API_BASE"] = settings.OLLAMA_BASE_URL
    
    # 打印新配置
    print(f"\n新 LiteLLM 配置:")
    print(f"- API 密钥: {litellm.api_key}")
    print(f"- Ollama API 基础 URL: {litellm.ollama_api_base}")
    print(f"- 模型列表: {json.dumps(litellm.model_list, indent=2) if hasattr(litellm, 'model_list') else '无'}")
    print(f"- 环境变量 OLLAMA_API_BASE: {os.environ.get('OLLAMA_API_BASE', '未设置')}")
    
    print("\n重置完成！")


if __name__ == "__main__":
    reset_litellm()